import re
import functools
import asyncio
import inspect
from concurrent.futures import ThreadPoolExecutor

def classify_sme_code(code):
    code = code.strip()  # 去除前后空格
    if re.match(r'^(600|601)\d{4}$', code):
        return "主板公司（上海证券交易所）"
    elif re.match(r'^(000)\d{4}$', code):
        return "主板公司（深圳证券交易所）"
    elif re.match(r'^(002)\d{4}$', code):
        return "中小企业板（深圳证券交易所）"
    elif re.match(r'^(300)\d{4}$', code):
        return "创业板（深圳证券交易所）"
    else:
        return "无效的股票代码"

# 示例
# stock_codes = ['600000', '000001', '002001', '300001', '688001', '01234']
# for code in stock_codes:
#     print(f"股票代码 {code} : {classify_sme_code(code)}")

class ThreadPoolCoroutineDecorator:
    def __init__(self, max_workers=None):
        self.max_workers = max_workers or (os.cpu_count() or 1) * 5
        self._pool = None

    def __call__(self, func):
        @functools.wraps(func)
        def wrapper(tasks, *args, **kwargs):
            task_chunks = self._chunk_tasks(tasks, self.max_workers)
            results = []
        
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = [
                    executor.submit(
                        self._run_task_chunk,
                        func,
                        chunk,
                        args,
                        kwargs,
                    ) for chunk in task_chunks
                ]

            for future in futures:
                results.append(future.result())

            return results

        return wrapper

    @staticmethod
    def _chunk_tasks(tasks, num_chunks):
        chunk_size = len(tasks) // num_chunks + (len(tasks) % num_chunks >0)
        return [tasks[i*chunk_size:(i+1)*chunk_size] for i in range(num_chunks)]

    def _run_task_chunk(self, func, task_chunk, args, kwargs):
        if inspect.iscoroutinefunction(func):
            return asyncio.run(self._run_async_tasks(func, task_chunk, args, kwargs))
        else:
            return self._run_sync_tasks(func, task_chunk, args, kwargs)

    @staticmethod
    def _run_sync_tasks(func, task_chunk, args, kwargs):
        results = []
        for task in task_chunk:
            try:
                result = await func(task, *args, **kwargs)
                results.append(result)
            except Exception as e:
                results.append((task, str(e)))
        return results

    @staticmethod
    def _run_async_tasks(func, task_chunk, args, kwargs):
        results = []
        for task in task_chunk:
            try:
                result = func(task, *args, **kwargs)
                results.append(result)
            except Exception as e:
                results.append((task, str(e)))
        return results




    

